package mq

import (
	"context"
	"time"

	"github.com/nsqio/go-nsq"
)

// MqHandlerFunc 定义消息处理函数格式
type MqHandlerFunc func(ctx context.Context, msg *nsq.Message) error

type middleware func(handler MqHandlerFunc) MqHandlerFunc

// NsqConsumer nsq consumer
type NsqConsumer struct {
	lookupds            []string
	lookupdPollInterval time.Duration
	maxInFlight         int
	consumers           []*nsq.Consumer
	middlewares         []middleware
}

// NewNsqConsumer new nsq consumer
func NewNsqConsumer(lookupds []string, pollInterval time.Duration, maxInFlight int) *NsqConsumer {
	return &NsqConsumer{
		lookupds:            lookupds,
		lookupdPollInterval: pollInterval,
		maxInFlight:         maxInFlight,
	}
}

// Close close consumer
func (n *NsqConsumer) Close() {
	for _, c := range n.consumers {
		c.Stop()
	}
}

// Use import middle ware
func (n *NsqConsumer) Use(mids ...middleware) {
	for _, mid := range mids {
		n.middlewares = append(n.middlewares, mid)
	}
}

func (n *NsqConsumer) wrap(h MqHandlerFunc) MqHandlerFunc {
	for i := len(n.middlewares) - 1; i >= 0; i-- {
		mid := n.middlewares[i]
		h = mid(h)
	}
	return h
}

func (n *NsqConsumer) toNsqHandler(h MqHandlerFunc) nsq.HandlerFunc {
	return func(msg *nsq.Message) error {
		return h(context.TODO(), msg)
	}
}

// Register register handler
func (n *NsqConsumer) Register(topic, channel string, handler MqHandlerFunc, concurrent int) error {
	cfg := nsq.NewConfig()
	cfg.LookupdPollInterval = n.lookupdPollInterval
	cfg.MaxInFlight = n.maxInFlight
	c, err := nsq.NewConsumer(topic, channel, cfg)
	if err != nil {
		return err
	}
	if concurrent == 0 {
		concurrent = 5
	}
	handler = n.wrap(handler)
	// 关闭系统日志
	c.SetLogger(nil, 0)
	c.AddConcurrentHandlers(n.toNsqHandler(handler), concurrent)
	n.consumers = append(n.consumers, c)
	return nil
}

// Start start consumer
func (n *NsqConsumer) Start() error {
	for _, c := range n.consumers {
		if err := c.ConnectToNSQLookupds(n.lookupds); err != nil {
			return err
		}
	}
	return nil
}
